#define HAVE_PTHREAD_H 1
#include <fstream>
#include "timer.h"
#include "lbfgs.h"
#include "common.h"
#include "DependencyParser2.h"
#include "DependencyEncoder.h"
#include "DependencyEncoder2OSibling.h"
#include "DependencyEncoder2OGrandChild.h"
#include "DependencyEncoder2O.h"
#include "DependencyEncoder3OGSibling.h"
#include "thread.h"

class ForestCreaterThread : public thread{
public:
	size_t start_i;
	size_t thread_num;
	size_t size;
	DependencyInstance** instList;
	DependencyPipe *pipe;
	void run(){
		char file[100];
		sprintf(file, "tmp/train%d.forest", start_i);
		ObjectWriter out = ObjectWriter(file);
		for (size_t i = start_i; i < size; i += thread_num) {
			DependencyInstance *inst = instList[i];
			inst->setFeatureVector(pipe->createFeatureVector(inst, true));
			pipe->writeInstance(inst, out);
			delete(inst);
		}
	}
};

class ModelTrainerThread : public thread{
public:
	size_t start_i;
	size_t thread_num;
	size_t numInst;
	int size;
	int *instLengths;
	DependencyPipe *pipe;
	ParserOptions *options;
	long double obj;
	double *param;
	std::vector<long double> expected;
	void run(){

		char file[100];
		sprintf(file, "tmp/train%d.forest", start_i);

		obj = 0.0;
		std::fill(expected.begin(), expected.end(), 0.0);
		int maxLength = 0;
		for(size_t i = start_i; i < numInst; i += thread_num){
			if(instLengths[i] > maxLength){
				maxLength = instLengths[i];
			}
		}
		
		DependencyEncoder *encoder = NULL;
		Parameters *parameter = new Parameters(size, param);
		if(options->order == 1){
			encoder = new DependencyEncoder(maxLength, pipe, file, parameter);
		}
		else if(options->order == 2){
			if(options->model == 1){
				encoder = new DependencyEncoder2OSibling(maxLength, pipe, file, parameter);
			}
			else if(options->model == 2){
				encoder = new DependencyEncoder2OGrandChild(maxLength, pipe, file, parameter);
			}
			else if(options->model == 3){
				encoder = new DependencyEncoder2O(maxLength, pipe, file, parameter);
			}
		}
		else if(options->order == 3){
			if(options->model == 2){
				encoder = new DependencyEncoder3OGSibling(maxLength, pipe, file, parameter);
			}
			else{
				cout<<"Model type error"<<endl;
				exit(0);
			}
		}
		else{
			cout<<"Model type error"<<endl;
			exit(0);
		}

		for(size_t i = start_i; i < numInst; i += thread_num){
			obj += encoder->gradient(&(expected[0]), instLengths[i]);
		}

		parameter->setParameters(NULL);
		delete(parameter);
		delete(encoder);
	}
};


DependencyParser2::DependencyParser2(DependencyPipe *pipe, ParserOptions *options){
	this->pipe = pipe;
	this->options = options;
}

DependencyParser2::~DependencyParser2(){
	delete(pipe);
}

void DependencyParser2::createForest(int thread_num, std::vector<DependencyInstance*> *instList){
	progress_timer pg;
	vector<ForestCreaterThread> thread(thread_num);
	for (size_t i = 0; i < thread_num; i++) {
		thread[i].start_i = i;
		thread[i].instList = &((*instList)[0]);
		thread[i].thread_num = thread_num;
		thread[i].size = instList->size();
		thread[i].pipe = pipe;
	}
	for (size_t i = 0; i < thread_num; i++) {
		thread[i].start();
	}

	for (size_t i = 0; i < thread_num; i++) {
		thread[i].join();
	}
	cout<<"Done!";
}

/*
bool DependencyParser2::trainModel(vector<int> *instLengths, double *param, int maxiter, int thread_num, double C, double eta, bool orthant){
	double old_obj = 1e+37;
	int    converge = 0;
	LBFGS lbfgs;
	std::vector<ModelTrainerThread> thread(thread_num);
	int size = pipe->dataAlphabet.size();
	size_t numInst = instLengths->size();

	for(size_t i = 0; i < thread_num; i++){
		thread[i].start_i = i;
		thread[i].size = size;
		thread[i].numInst = numInst;
		thread[i].thread_num = thread_num;
		thread[i].instLengths = &((*instLengths)[0]);
		thread[i].expected.resize(size);
		thread[i].pipe = pipe;
		thread[i].options = options;
		thread[i].param = param;
	}

	for (size_t itr = 0; itr < maxiter; ++itr) {
		fflush(stdout);
		progress_timer pg;
		for (size_t i = 0; i < thread_num; ++i) thread[i].start();
		for (size_t i = 0; i < thread_num; ++i) thread[i].join();

		for (size_t i = 1; i < thread_num; ++i) {
			thread[0].obj += thread[i].obj;
		}

		for (size_t i = 1; i < thread_num; ++i) {
			for (int j = 0; j < size; ++j) {
				thread[0].expected[j] += thread[i].expected[j];
			}
		}

		size_t num_nonzero = 0;
		if (orthant) {   // L1
			 for (int j = 0; j < size; ++j) {
				 thread[0].obj += std::fabs(param[j] / C);
				if (param[j] != 0.0) ++num_nonzero;
			}
		} 
		else{
			num_nonzero = size;
			for (int j = 0; j < size; ++j) {
				thread[0].obj += (param[j] * param[j] /(2.0 * C));
				thread[0].expected[j] += param[j] / C;
			}
		}

		double diff = (itr == 0 ? 1.0 : std::abs(old_obj - thread[0].obj)/old_obj);
		printf("iter=%d act=%d obj=%.8lf diff=%.8lf time=", itr, num_nonzero, thread[0].obj, diff);
		old_obj = thread[0].obj;

		if (diff < eta){
			converge++;
		}
		else{
			converge = 0;
		}

		if (itr > maxiter || converge == 3){
			break;  // 3 is ad-hoc
		}

		
		if (lbfgs.optimize(size, param, thread[0].obj, &(thread[0].expected[0]), orthant, C) <= 0){
			return false;
		}
		
  }

  return true;
}
*/

bool DependencyParser2::trainModel(vector<int> *instLengths, double *param, int maxiter, int thread_num, double C, double eta, bool orthant){
	long double old_obj = 1e+37;
	int converge = 0;
	int converge1 = 0, converge2 = 0, converge3 = 0;
	double eta1 = 0.0001, eta2 = 0.00001, eta3 = 0.000001;
	bool flag1 = false, flag2 = false, flag3 = false;
	LBFGS lbfgs;

	//init threads
	std::vector<ModelTrainerThread> thread(thread_num);
	int size = pipe->dataAlphabet.size();
	size_t numInst = instLengths->size();

	for(size_t i = 0; i < thread_num; i++){
		thread[i].start_i = i;
		thread[i].size = size;
		thread[i].numInst = numInst;
		thread[i].thread_num = thread_num;
		thread[i].instLengths = &((*instLengths)[0]);
		thread[i].expected.resize(size);
		thread[i].pipe = pipe;
		thread[i].options = options;
		thread[i].param = param;
	}
	
	//gradient
	double *g = new double[size];

	for (size_t itr = 0; itr < maxiter; ++itr) {
		fflush(stdout);
		progress_timer pg;
		for (size_t i = 0; i < thread_num; ++i) thread[i].start();
		for (size_t i = 0; i < thread_num; ++i) thread[i].join();

		for (size_t i = 1; i < thread_num; ++i) {
			thread[0].obj += thread[i].obj;
		}

		for (size_t i = 1; i < thread_num; ++i) {
			for (int j = 0; j < size; ++j) {
				thread[0].expected[j] += thread[i].expected[j];
			}
		}

		size_t num_nonzero = 0;
		if (orthant) {   // L1
			 for (int j = 0; j < size; ++j) {
				 thread[0].obj += std::fabs(param[j] / C);
				if (param[j] != 0.0) ++num_nonzero;
			}
		} 
		else{
			num_nonzero = size;
			for (int j = 0; j < size; ++j) {
				thread[0].obj += (param[j] * param[j] /(2.0 * C));
				thread[0].expected[j] += param[j] / C;
			}
		}

		//get gradient
		for(int j = 0; j < size; j++){
			g[j] = thread[0].expected[j];
		}

		long double diff = (itr == 0 ? 1.0 : std::abs(old_obj - thread[0].obj)/old_obj);
		printf("iter=%d act=%d obj=%.8Lf diff=%.8Lf time=", itr, num_nonzero, thread[0].obj, diff);
		old_obj = thread[0].obj;

		if(diff < eta1 && !flag1){
			converge1++;
		}
		else{
			converge1 = 0;
		}
		if(!flag1 && converge1 == 3){
			string str = string("parse_model-0.0001.pm");
			printf("\n Saving Model %.8lf ...", eta1);
			saveModel(&str, param);
			printf("Done\n");
			flag1 = true;
		}

		if(diff < eta2 && !flag2){
			converge2++;
		}
		else{
			converge2 = 0;
		}
		if(!flag2 && converge2 == 3){
			string str = string("parse_model-0.00001.pm");
			printf("\n Saving Model %.8lf ...", eta2);
			saveModel(&str, param);
			printf("Done\n");
			flag2 = true;
		}

		if(diff < eta3 && !flag3){
			converge3++;
		}
		else{
			converge3 = 0;
		}
		if(!flag3 && converge3 == 3){
			string str = string("parse_model-0.000001.pm");
			printf("\n Saving Model %.8lf ...", eta3);
			saveModel(&str, param);
			printf("Done\n");
			flag3 = true;
		}

		if (diff < eta){
			converge++;
		}
		else{
			converge = 0;
		}

		if (itr > maxiter || converge == 3){
			break;  // 3 is ad-hoc
		}

		
		if (lbfgs.optimize(size, param, thread[0].obj, g, orthant, C) <= 0){
			return false;
		}
	}

	delete[](g);
	return true;
}

void DependencyParser2::saveModel(std::string *file, double *param){
	ObjectWriter out(file->c_str());
	int size = (int)(pipe->dataAlphabet.size());
	out.writeArray(param, size);
	pipe->dataAlphabet.writeObject(out);
	pipe->typeAlphabet.writeObject(out);
}

bool DependencyParser2::learn(string *train_file, string *model_file, double C, double eta, int maxiter, int thread_num, int algorithm){
	CHECK_FALSE(eta > 0.0) << "eta must be > 0.0";
	CHECK_FALSE(C >= 0.0) << "C must be >= 0.0";
	CHECK_FALSE(thread_num > 0) << "thread must be > 0";

#ifndef CRFPP_USE_THREAD
	CHECK_FALSE(thread_num == 1)
      << "This architecture doesn't support multi-thrading";
#endif

#define WHAT_ERROR(msg) do {                                    \
    delete(instLengths);                                        \
    std::cerr << msg << std::endl;                              \
    return false; } while (0)


	vector<DependencyInstance*> instList(0);
	vector<int> *instLengths = pipe->createInstances(train_file, &instList);
	if(options->createForest){
		createForest(thread_num, &instList);
	}

	printf("Number of sentences:   %d\n", instLengths->size());
	printf("Number of features:    %d\n", pipe->dataAlphabet.size());
	printf("Number of edge labels: %d\n", pipe->typeAlphabet.size());
	printf("Number of thread(s):   %d\n", thread_num);
	printf("eta:                   %.8lf\n", eta);
	printf("C:                     %.8lf\n", C);
	fflush(stdout);

	progress_timer pg;

	int size_param = (int)(pipe->dataAlphabet.size());
	double *param = new double[size_param];
	for(int i = 0; i < size_param; i++){
		param[i] = 0.0;
	}

	switch(algorithm){
		case 2:
			if(!trainModel(instLengths, param, maxiter, thread_num, C, eta, false)){
				WHAT_ERROR("LBFGS-L2 execute error");
			}
			break;
		case 3:
			if(!trainModel(instLengths, param, maxiter, thread_num, C, eta, true)){
				WHAT_ERROR("LBFGS-L1 execute error");
			}
			break;
	}

	cout<<"Saving model...";
	saveModel(model_file, param);
	cout<<"Done"<<endl;
	delete(instLengths);
	delete[](param);

	cout<<"\nDone!";
	return true;
}
